#!/usr/bin/python
# -*- coding: UTF-8 -*-


import psycopg2
import pandas as pd
from psycopg2 import sql
from psycopg2.extras import RealDictCursor  
# from io import StringIO 
import geopandas as gpd  
from shapely.geometry import Point 
from shapely.wkb import dumps
# from sqlalchemy import create_engine
# import geoalchemy2
  
# 数据库连接参数  
conn_params = {  
  'dbname': 'postgis4sew',    # 数据库名  
  'user': 'postgres',         # 数据库用户名  
  'password': 'Postgres~0413',# 数据库密码  
  'host': '100.88.24.8',      # 数据库主机地址，通常是'localhost'或IP地址  
  'port': '5433'              # 数据库端口，默认是5432  
}


def read_csv_to_dataframe(file_path, encoding='utf-8'):
  """
  从CSV文件读取数据并转换为DataFrame对象
  :param file_path: CSV文件的路径
  :param encoding: 文件编码格式，默认为'utf-8'
  :return: 转换后的DataFrame对象
  """
  data_raw = pd.read_csv(file_path, encoding=encoding)
  df = data_raw.loc[:, ['部件标识码', '坐标X', '坐标Y', '类别', '所在道路名称',
                                  '自然路面标高', '井深', '井底标高', '井盖形状', '井盖尺寸1', '井盖尺寸2']]
  df.rename(columns={'部件标识码': 'name', '坐标X': 'x', '坐标Y': 'y', '类别': 'type',
                               '所在道路名称': 'road_name', '自然路面标高': 'ground_elevation',
                               '井深': 'depth', '井底标高': 'bottom_elevation',
                               '井盖形状': 'cover_shape', '井盖尺寸1': 'cover_size_1', '井盖尺寸2': 'cover_size_2'},
                               inplace=True)
  df.dtypes
  df['name'] = df['name'].astype('string')
  df['x'] = df['x'].astype('float64')
  df['y'] = df['y'].astype('float64')
  df['type'] = df['type'].astype('string')
  df['road_name'] = df['road_name'].astype('string')
  df['road_name'] = df['road_name'].fillna("")
  df['ground_elevation'] = df['ground_elevation'].astype('float64')
  df['depth'] = df['depth'].astype('float')
  df['bottom_elevation'] = df['bottom_elevation'].astype('float64')
  df['cover_shape'] = df['cover_shape'].astype('string')
  df['cover_shape'] = df['cover_shape'].fillna("")
  df['cover_size_1'] = df['cover_size_1'].astype('float')
  df['cover_size_2'] = df['cover_size_2'].astype('float')
  df['cover_size_1']=df['cover_size_1'].fillna(0)
  df['cover_size_2']=df['cover_size_2'].fillna(0)

  return df


def df_to_gdf(dataframe, crs='EPSG:4326'):
  """
  将Pandas DataFrame转换为GeoDataFrame对象
  :param dataframe: 要转换的DataFrame对象
  :return: 转换后的GeoDataFrame对象
  """
  geometry = [Point(x, y) for x, y in zip(dataframe['x'], dataframe['y'])]
  gdf = gpd.GeoDataFrame(dataframe, geometry=geometry, crs=crs)
  return gdf


# def add_data_to_table(cursor, table_name, dataframe):
#   """
#   将Pandas DataFrame插入到PostgreSQL表中。  
  
#   参数:  
#     cursor (psycopg2.Cursor): 数据库游标对象。  
#     table_name (str): 目标表名。  
#     dataframe (pd.DataFrame): 要插入的DataFrame。  
  
#   返回:  
#     None  
#   """
#   # 构造INSERT语句  
#   placeholders = ', '.join(['%s'] * len(dataframe.columns))  
#   cols = ', '.join(dataframe.columns)  
#   insert_sql = f"INSERT INTO {table_name} ({cols}) VALUES ({placeholders})"  
  
#   # 将DataFrame转换为记录列表，其中每条记录都是一个元组  
#   # records = dataframe.to_records(index=False, column_names=True).tolist()  
#   # 将DataFrame转换为记录数组，然后转换为元组列表  
#   records_array = dataframe.to_records(index=False)  
#   records_list = [tuple(record) for record in records_array] 
#   # 插入数据
#   cursor.executemany(insert_sql, records_list)


def add_few_geodata_to_node_table(cur, table_name, gdf, srid):
  """
  将GeoDataFrame插入到PostgreSQL表中。  

  参数:  
    cursor (psycopg2.Cursor): 数据库游标对象。  
    table_name (str): 目标表名。  
    dataframe (pd.DataFrame): 要插入的GeoDataFrame。
  """
  # 构造插入语句
  insert_sql = sql.SQL("""
    INSERT INTO {tableName} (name, x, y, type, road_name, ground_elevation, depth, bottom_elevation,
                            cover_shape, cover_size_1, cover_size_2, geometry)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, ST_GeomFromText(%s, %s))
    """
    ).format(
      tableName=sql.Identifier(table_name)
    )
  for _, row in gdf.iterrows():
    df = tuple(row.drop(['geometry']).tolist()) + (row['geometry'].wkt,) + (srid,)
    cur.execute(insert_sql, df)


def create_geomtry_table(cursor, table_name):
  """
  创建包含几何字段的表
  :param cursor: 数据库游标
  :tableName: 表名
  :param geomType:几何类型, 如'POINT', 'LINESTRING', 'POLYGON'.
                  格式为字符串，如'GEOMETRY(Point, 4326)'
          id SERIAL PRIMARY KEY,
          name character varying(255) NOT NULL,
          geom {geometry_type}
  """
  # try:
  # 创建表的SQL语句
  create_table_sql = sql.SQL("""
      CREATE TABLE IF NOT EXISTS {table_name}();
  """).format(
      table_name=sql.Identifier(table_name)
      # table_name=sql.Identifier(table_name),
      # geometry_type=sql.Identifier(geom_type)
  )

  # 执行SQL语句
  cursor.execute(create_table_sql)

  #   print(f"Table {table_name} created successfully with geometry type {geom_type}.")
  # except Exception as e:
  #   print(f"An error occurred: {e}")


def add_columns_to_table(cur, table_name, columns):  
  """  
  向指定的表中添加新列。  
  
  :param cur: psycopg2数据库连接游标
  :param tableName: 要添加新列的表名  
  :param columns: 包含新列名和数据类型的列表，每个元素是一个(列名, 数据类型)的元组  
  """  
  # 示例用法（注意：这里没有包括数据库连接的代码）  
  # cursor = ...  # 获取游标对象  
  # table_name = 'your_table'  
  # columns_to_add = [('new_column1', 'INTEGER'), ('new_column2', 'VARCHAR(100)')]  
  # add_columns_to_table(cursor, table_name, columns_to_add)
  # 遍历列名列表，为每个新列构建ALTER TABLE语句
  for column_name, data_type in columns:
    add_column_query = sql.SQL("""
      ALTER TABLE {table}
      ADD COLUMN {column} {type}
      """).format(
        table=sql.Identifier(table_name),  
        column=sql.Identifier(column_name),  
        type=sql.SQL(data_type)  
    )
    # 执行ALTER TABLE语句  
    cur.execute(add_column_query)


def get_tables(cursor):  
  """  
  列出数据库中的所有表。  
 
  :param cursor: psycopg2游标对象  
  """  
  try:  
      # 查询信息架构中的表名  
      query = sql.SQL("SELECT tablename FROM pg_tables WHERE schemaname = 'public';")  
      cursor.execute(query)  
 
      # 获取所有表的名称
      tables = cursor.fetchall()
 
      # 打印表名  
      tables1 = [element for sublist in tables for element in sublist]
      for table in tables:  
          print(table[0])  
      return tables1

  except (Exception, psycopg2.DatabaseError) as error:  
      print(f"Error: {error}")  


def get_table_columns(conn, table_name, schema_name='public'):  
  """  
  查询指定表的列名。  

  :param conn: 数据库连接对象  
  :param table_name: 表名  
  :param schema_name: 模式名(默认为'public')
  :return: 列名的列表  
  """  
  try:  
    with conn.cursor() as cursor:  
      # 使用参数化查询防止SQL注入  
      query = sql.SQL("""  
        SELECT column_name  
        FROM information_schema.columns  
        WHERE table_name = %s  
          AND table_schema = %s;  
      """)  
      cursor.execute(query, (table_name, schema_name))  
      columns = cursor.fetchall()  
        
      # 提取列名  
      column_names = [column[0] for column in columns]  
        
      return column_names  
  except (Exception, psycopg2.DatabaseError) as error:  
    print(f"Error: {error}")  
    return []  


if __name__ == "__main__":

  # 读取预处理数据
  file_name='manhole_all.csv'
  manhole_data = read_csv_to_dataframe(file_name)
  manhole_data_drop_dup = manhole_data.drop_duplicates(subset=['name'], keep=False)
  # 创建gdf
  manhole_data_gdf = df_to_gdf(manhole_data, crs='EPSG:4490')

  # sql建立链接
  # engine = create_engine("postgresql://postgres:Postgres~0413@100.88.24.8:5433/postgis4sew")
  # 写入数据
  # manhole_data_gdf.to_postgis('manhole', engine, if_exists='append', index=False)
  # 建立连接  
  try:  
    conn = psycopg2.connect(**conn_params)  
    print("连接成功！")  
  except Exception as e:  
    print(f"连接失败：{e}")  
    conn = None
  
  # 执行操作
  if conn:
    cur = conn.cursor()
    # 查询表
    tables = get_tables(cur)
    if ( 'manhole' not in tables) :
      # 如不存在manhole表，则创建
      try:
        create_geomtry_table(cur, 'manhole')
        # 向表中添加列
        columns_to_add = [('name', 'VARCHAR(50)'), ('type', 'CHAR(10)'),
                          ('x', 'DOUBLE PRECISION'), ('y', 'DOUBLE PRECISION'),
                          ('depth', 'REAL'),
                          ('ground_elevation', 'DOUBLE PRECISION'), ('bottom_elevation', 'DOUBLE PRECISION'),
                          ('cover_shape', 'VARCHAR(20)'), ('cover_size_1', 'REAL'), ('cover_size_2', 'REAL'),
                          ('road_name', 'VARCHAR(50)'),
                          ('geometry', 'GEOMETRY')]
        add_columns_to_table(cur, 'manhole', columns_to_add)
        conn.commit()  
      except (Exception, psycopg2.DatabaseError) as error:  
        print(f"创建表失败：{error}")
        conn.rollback()

  # 插入数据
  if conn:
    if ( 'manhole' in tables):
      try:
        add_few_geodata_to_node_table(cur, 'manhole', manhole_data_gdf, 4490)
        conn.commit()  
      except (Exception, psycopg2.DatabaseError) as error:  
        print(f"插入数据失败：{error}")
        conn.rollback()
 
  if conn:
    if cur:
      cur.close()
    conn.close()


